package com.eudon.ai.agent.etl.extension;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
import org.springframework.stereotype.Component;

import java.util.List;

/**
 * @author eudon
 * @description 查询扩展器，当用户查询进来时，调用此扩展器可以将用户的问题扩展为多个更专业的问题
 * @date 2025/8/15
 * @email syd19970616@gmail.com
 * @address <a href="https://gitee.com/eudon">Gitee</a>
 */
@Component
public class MyMultiQueryExpander {
    private final ChatClient.Builder chatClientBuilder;
    public MyMultiQueryExpander(ChatModel dashscopeChatModel) {
        this.chatClientBuilder = ChatClient.builder(dashscopeChatModel);
    }

    public List<Query> expand(Query query,int queryNumber,boolean includeOriginal) {
        MultiQueryExpander multiQueryExpander = MultiQueryExpander.builder().numberOfQueries(queryNumber)
                .chatClientBuilder(chatClientBuilder)
                .includeOriginal(includeOriginal)
                .build();
        return multiQueryExpander.expand(query);

    }
}
